import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D


def scatter_3d(X=None, x=None, y=None, z=None):
    if X is not None:
        x = X[0]
        y = X[1]
        z = X[2]
    fig = plt.figure()  # 定义新的三维坐标轴
    ax = Axes3D(fig)
    ax.scatter3D(x, y, z, cmap='Blues')  # 绘制散点图
    plt.show()


def plot_3d(X=None, x=None, y=None, z=None):
    if X is not None:
        x = X[0]
        y = X[1]
        z = X[2]
    fig = plt.figure()  # 定义新的三维坐标轴
    ax = Axes3D(fig)
    ax.plot3D(x, y, z, 'gray')  # 绘制空间曲线
    plt.show()


def plot_surface_3d(X=None, x=None, y=None, z=None, rstride=10, cstride=10):
    """
        plot_surface: 生成平面
        contour: 等高线(即投影图)
        contourf
    """
    if X is not None:
        x = X[0]
        y = X[1]
        z = X[2]
    # 定义新的三维坐标轴
    plt.figure()
    ax = plt.axes(projection='3d')
    ax.plot_surface(x, y, z, rstride=rstride, cstride=cstride, cmap='rainbow')  # 如果加入渲染时的步长, 会得到更加清晰细腻的图像. 两个值越小越细腻
    ax.contour(x, y, z, offset=-2, cmap='rainbow')  # 等高线图，要设置offset，为Z的最小值
    plt.show()


def plot_surface2_3d(X=None, x=None, y=None, z=None, rstride=10, cstride=10):
    """
        plot_surface: 生成平面
        contour: 等高线(即投影图)
        contourf: 等高面(即投影填充)
        set_xlabel: 设定x轴显示标签
        set_xlim: 设定x轴显示范围
    """
    if X is not None:
        x = X[0]
        y = X[1]
        z = X[2]
    # 定义坐标轴
    fig3 = plt.figure()
    ax3 = plt.axes(projection='3d')
    # 作图
    ax3.plot_surface(x, y, z, alpha=0.3, cmap='winter')  # 生成表面， alpha 用于控制透明度
    ax3.contour(x, y, z, zdir='z', offset=-3, cmap="rainbow")  # 生成z方向投影，投到x-y平面
    ax3.contour(x, y, z, zdir='x', offset=-6, cmap="rainbow")  # 生成x方向投影，投到y-z平面
    ax3.contour(x, y, z, zdir='y', offset=6, cmap="rainbow")  # 生成y方向投影，投到x-z平面
    # ax3.contourf(x, y, z, zdir='y', offset=6, cmap="rainbow")  # 生成y方向投影填充，投到x-z平面，contourf()函数
    # 设定显示范围
    ax3.set_xlabel('X')
    ax3.set_xlim(-6, 4)  # 拉开坐标轴范围显示投影
    ax3.set_ylabel('Y')
    ax3.set_ylim(-4, 6)
    ax3.set_zlabel('Z')
    ax3.set_zlim(-3, 3)

    plt.show()


# %%
# 将降维后的数据可视化,2维

def plot_embedding_2d(X, y, title=None):
    # 坐标缩放到[0,1]区间
    # x_min, x_max = np.min(X, axis=0), np.max(X, axis=0)
    # X = (X - x_min) / (x_max - x_min)
    # 降维后的坐标为（X[i, 0], X[i, 1]），在该位置画出对应的digits
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    for i in range(X.shape[0]):
        ax.text(X[i, 0], X[i, 1], str(y[i]),
                color=plt.cm.Set1(y[i]),
                fontdict={'weight': 'bold', 'size': 9})
    if title is not None:
        plt.title(title)
    ax.set_xlim(-5, 5)
    ax.set_ylim(-5, 5)
    plt.show()


# %%
# 将3维数据可视化
def plot_embedding_3d(X, y, title=None, flag=False):
    """
    :param title:
    :param flag: 是否抽样显示
    :return:
    """
    # 坐标缩放到[0,1]区间
    # x_min, x_max = np.min(X, axis=0), np.max(X, axis=0)
    # X = (X - x_min) / (x_max - x_min)
    # 降维后的坐标为（X[i, 0], X[i, 1],X[i,2]），在该位置画出对应的digits
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1, projection='3d')
    for i in range(X.shape[0]):
        if flag & (i & 1 == 0):  # 为True并且i为偶数
            continue
        ax.text(X[i, 0], X[i, 1], X[i, 2], str(y[i]), color=plt.cm.Set1(y[i])
                , fontdict={'weight': 'bold', 'size': 9})
        if title is not None:
            plt.title(title)
    ax.set_xlabel('X')
    ax.set_xlim(-3, 3)  # 拉开坐标轴范围显示投影
    ax.set_ylabel('Y')
    ax.set_ylim(-3, 3)
    ax.set_zlabel('Z')
    ax.set_zlim(-3, 3)
    plt.show()
